package middleware

import (
	"bytes"
	"context"
	"github.com/gin-gonic/gin"
	"io"
	"time"
)

type LogMiddlewareBuilder struct {
	// logFunc 日志函数
	logFunc       func(ctx context.Context, al AccessLog)
	allowReqBody  bool // 是否允许记录请求体，避免黑客恶意发起请求
	allowRespBody bool // 是否允许记录响应体
}

func NewLogMiddlewareBuilder(logFunc func(ctx context.Context, al AccessLog)) *LogMiddlewareBuilder {
	return &LogMiddlewareBuilder{
		logFunc: logFunc,
	}
}

// AccessLog 日志结构体
// 具体可以根据公司的规范来设计
type AccessLog struct {
	Path       string        `json:"path"`        // 请求路径
	Method     string        `json:"method"`      // 请求方法
	ReqBody    string        `json:"req_body"`    // 请求体
	RespBody   string        `json:"resp_body"`   // 响应体
	Duration   time.Duration `json:"duration"`    // 耗时
	StatusCode int           `json:"status_code"` // 状态码
}

func (l *LogMiddlewareBuilder) AllowReqBody() *LogMiddlewareBuilder {
	l.allowReqBody = true
	return l
}

func (l *LogMiddlewareBuilder) AllowRespBody() *LogMiddlewareBuilder {
	l.allowRespBody = true
	return l
}

// Build 构建日志中间件
func (l *LogMiddlewareBuilder) Build() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		path := ctx.Request.URL.Path
		// 有些黑客跟你过不去，它的path可能很长，你可以考虑截取一部分
		if len(path) > 1024 { // todo: 这个1024你可以做一个字段
			path = path[:1024]
		}
		method := ctx.Request.Method
		al := AccessLog{
			Path:   path,
			Method: method,
		}
		if l.allowRespBody {
			// 这里要注意，request.Body 是一个Stream对象，只能读取一次，
			// 所以需要再给它放回去，不然后续的步骤是读不到的
			body, _ := ctx.GetRawData()
			// 防止body过大，截取一部分
			if len(body) > 2048 {
				body = body[:2048]
			} else {
				al.ReqBody = string(body)
			}
			// 把body放回去
			ctx.Request.Body = io.NopCloser(bytes.NewReader(body))
		}

		start := time.Now()
		// 如果允许记录响应体，替换为我们自定义的responseWriter
		if l.allowRespBody {
			ctx.Writer = &responseWriter{
				ResponseWriter: ctx.Writer,
				al:             &al,
			}
		}
		// 就是说，下次回到这里的时候，已经执行完了业务逻辑，即可以拿到响应体了
		defer func() {
			al.Duration = time.Since(start)
			if l.allowRespBody {
				// ctx 是没有 Response 的，那么我该怎么拿到响应体呢？
				// 我们就来一个骚操作，他原本的 ResponseWriter 中没有想用的 Body，
				// 那么我们就给利用装饰器自己定义一个
				l.logFunc(ctx, al) // 到此，AccessLog 记录完毕，打印日志
			}
		}()

		// 直接执行下一个中间件... 直到业务逻辑
		ctx.Next()
	}
}

// 利用装饰器对 ResponseWriter 进行包装，达到记录响应体的目的
type responseWriter struct {
	al *AccessLog
	gin.ResponseWriter
}

// Write 覆盖原本的Write方法
func (w *responseWriter) Write(data []byte) (int, error) {
	w.al.RespBody = string(data)
	return w.ResponseWriter.Write(data)
}

// WriteHeader 覆盖原本的WriteHeader方法
func (w *responseWriter) WriteHeader(statusCode int) {
	w.al.StatusCode = statusCode
	w.ResponseWriter.WriteHeader(statusCode)
}
